from utils import *
import matplotlib.pyplot as plt
import os
from noanchor_qa import gen_noanchor_qa
from anchor_qa import gen_anchor_qa
from extra1_qa import gen_extra1_qa
from extra2_qa import gen_extra2_qa
from config import *

color_info = read_json(COLOR_PATH)

assets_info = read_json(ASSETS_PATH)
old_assets_info = read_json(OLD_ASSETS_PATH)

path = os.path.expanduser('~/scenegen/data_extra2')
# id = '0ba87c9b13ba8562'
id = 'fe1c96a01856d939'
json_path = os.path.join(path, id + '.json')
depth_path = os.path.join(path, id + '.png')
scene_json = read_json(json_path)

scene_type = 'common' if 'type' not in scene_json else scene_json['type']
cam = scene_json['camera']
table = scene_json['table']
obj_dict = scene_json['objects']
tablez = table['tablemax'][2]
table = (table['tablemin'][0], table['tablemin'][1], table['tablemax'][0], table['tablemax'][1], tablez)

# obj0 = objs[0]
# mask = mask_decoder(obj0['mask'])

# plt.matshow(mask, cmap='Greys')
# plt.title("OBJECT0 MASK")
# plt.axis('off')
# plt.show()

# print('center of object0:', get_mask_center(mask))

# process basic info
for id, info in obj_dict.items():
    id = get_original_id(id)
    info['id'] = id
    info['color'] = color_info[id]
    info['mask'] = mask_decoder(info['mask'])
    info['visible'] = info['mask'].any()
    info['3dloc'] = ((info['bboxmax'][0] + info['bboxmin'][0]) / 2, (info['bboxmax'][1] + info['bboxmin'][1]) / 2, (info['bboxmax'][2] + info['bboxmin'][2]) / 2)
    info['3dtabloc'] = (info['3dloc'][0], info['3dloc'][1], tablez)
    if info['visible']:
        info['2dloc'] = get_mask_center(info['mask'])
    else:
        info['2dloc'] = None
    info['2dtabloc'] = world_to_camera_pixel(info['3dtabloc'], cam)
    info['visible'] = info['visible'] and check_enough_visible_part(info, cam)
    info['bottom_visible'] = check_bottom_visible(info, cam)
    info['bottom_points'] = [
        [info['3dloc'][0], info['bboxmin'][1], tablez],
        [info['3dloc'][0], info['bboxmax'][1], tablez],
        [info['bboxmin'][0], info['3dloc'][1], tablez],
        [info['bboxmax'][0], info['3dloc'][1], tablez]
    ]

    dir_desc = info['tempalte_dir']
    if scene_type != 'freeorient':
        if id not in assets_info:
            info['direction'] = [0, 0, 1]
        else:
            new_init_direction = assets_info[id]['direction']
            if new_init_direction[2] != 0:
                info['direction'] = [0, 0, 1]
            else:
                dir_desc = assets_info[id]['tempalte_dir']
                old_init_direction = old_assets_info[id]['direction']
                old_direction = info['direction']
                old_init_angle = get_angle(old_init_direction[0], old_init_direction[1])
                old_angle = get_angle(old_direction[0], old_direction[1])
                new_init_angle = get_angle(new_init_direction[0], new_init_direction[1])
                new_angle = new_init_angle + (old_angle - old_init_angle)
                info['direction'] = [np.cos(new_angle), np.sin(new_angle), 0]
    info['dir_desc'] = dir_desc
    info['exclusive_ref'] = []


if not check_nomask(obj_dict): # over 2 invisible masks
    exit(0)
if not check_depth(depth_path): # didn't render completely
    exit(0)

# get exclusive description
tmp_objs = {}
for id, info in obj_dict.items():
    if not info['visible']:
        continue
    tmp_objs[id] = info
calc_rank(obj_dict, tmp_objs, 'object', '')

cates_set = set()
cates_count = {}
for id, info in obj_dict.items():
    cates_set.add(info['cate'])
    cates_count[info['cate']] = cates_count.get(info['cate'], 0) + 1
for cate in cates_set:
    cate_obj = {}
    name = None
    for id, info in obj_dict.items():
        if not info['visible']:
            continue
        if info['cate'] == cate:
            cate_obj[id] = info
            name = info['minimal_desc']
    calc_rank(obj_dict, cate_obj, name, '')


color_cates_count = {}
if scene_type == 'color':
    color_cates_set = set()
    for id, info in obj_dict.items():
        if info['color'] != 'none':
            color_cates_set.add((info['cate'], info['color']))
            color_cates_count[(info['cate'], info['color'])] = color_cates_count.get((info['cate'], info['color']), 0) + 1
    for (cate, color) in color_cates_set:
        cate_obj = {}
        name = None
        for id, info in obj_dict.items():
            if not info['visible']:
                continue
            if info['cate'] == cate and info['color'] == color:
                cate_obj[id] = info
                name = info['minimal_desc']
        calc_rank(obj_dict, cate_obj, name, color)

for id, info in obj_dict.items():
    print(id)
    print(info['color'])
    print(info['minimal_desc'])
    print(info['2dloc'])
    print(info['2dtabloc'])
    print(info['bottom_visible'])
    print(info['exclusive_ref'])

basic_qa = []

# no anchor qa
basic_qa.extend(gen_noanchor_qa(obj_dict))

# anchor qa
basic_qa.extend(gen_anchor_qa(obj_dict, cam, table, cates_count, color_cates_count))

extra1_qa = []
if scene_type == 'equidistance':
    extra1_qa = gen_extra1_qa(obj_dict, cam, table)

extra2_qa = []
if scene_type == 'freeorient':
    extra2_qa = gen_extra2_qa(obj_dict, cam, table)
print(extra2_qa)
